import tkinter as tk
import random
import math

# 常量定义
WIDTH, HEIGHT = 500, 500
NUM_OBJECTS = 300
MIN_DISTANCE = 2  # 物体之间的最小距离
STEP_SIZE = 10  # 每次移动的步长

# 颜色映射（使用Tkinter支持的颜色名称）
DIRECTION_COLORS = {
    "up": "red",     # 上：红色
    "down": "orange",# 下：橙色
    "left": "blue",# 左：黄色
    "right": "green" # 右：绿色
}

class Object:
    def __init__(self, canvas):
        self.canvas = canvas
        self.x = random.randint(0, WIDTH-1)
        self.y = random.randint(0, HEIGHT-1)
        self.color = "black"
        self.moving = False
        self.angle = 0  # 初始角度（后续会被覆盖）
        self.id = canvas.create_rectangle(self.x, self.y, self.x+1, self.y+1, fill=self.color, outline=self.color)

    def move(self):
        if self.moving:
            # 计算新位置
            rad = math.radians(self.angle)
            dx = STEP_SIZE * math.cos(rad)
            dy = STEP_SIZE * math.sin(rad)
            new_x = self.x + dx
            new_y = self.y + dy

            # 边界碰撞检测
            hit_wall = False
            if new_x < 0 or new_x >= WIDTH or new_y < 0 or new_y >= HEIGHT:
                self.bounce()
                hit_wall = True

            # 撞墙后重新计算位置
            if hit_wall:
                rad = math.radians(self.angle)
                dx = STEP_SIZE * math.cos(rad)
                dy = STEP_SIZE * math.sin(rad)
                new_x = self.x + dx
                new_y = self.y + dy

            # 更新位置
            self.x, self.y = new_x, new_y
            self.canvas.coords(self.id, self.x, self.y, self.x+1, self.y+1)
            
            # 强制更新颜色
            self.update_color()

    def bounce(self):
        """光线反射逻辑"""
        # 计算当前运动方向
        rad = math.radians(self.angle)
        dx = STEP_SIZE * math.cos(rad)
        dy = STEP_SIZE * math.sin(rad)

        # 水平碰撞（左右墙）
        if self.x + dx < 0 or self.x + dx >= WIDTH:
            self.angle = 180 - self.angle
        
        # 垂直碰撞（上下墙）
        if self.y + dy < 0 or self.y + dy >= HEIGHT:
            self.angle = -self.angle
        
        # 标准化角度到0-360
        self.angle %= 360
        self.update_color()  # 碰撞后立即更新颜色

    def update_color(self):
        """根据运动方向更新颜色"""
        if 45 <= self.angle < 135:     # 向上运动
            color = DIRECTION_COLORS["up"]
        elif 135 <= self.angle < 225:  # 向左运动
            color = DIRECTION_COLORS["left"]
        elif 225 <= self.angle < 315:  # 向下运动
            color = DIRECTION_COLORS["down"]
        else:                          # 向右运动
            color = DIRECTION_COLORS["right"]
        
        if self.color != color:
            self.color = color
            # 显式更新画布上的颜色
            self.canvas.itemconfig(self.id, fill=color, outline=color)

    def set_moving(self, moving, angle=None):
        self.moving = moving
        if moving:
            if angle is not None:
                self.angle = angle
            else:
                # 初始随机方向（确保是四个主要方向之一）
                self.angle = random.choice([0, 90, 180, 270])  # 右/上/左/下
            self.update_color()

class EvolutionSimulation:
    def __init__(self, root):
        self.root = root
        self.canvas = tk.Canvas(root, width=WIDTH, height=HEIGHT, bg="white")
        self.canvas.pack(side=tk.LEFT)

        # 初始化物体（全部静止）
        self.objects = [Object(self.canvas) for _ in range(NUM_OBJECTS)]
        self.moving_objects = []

        # 控制按钮
        control_frame = tk.Frame(root)
        control_frame.pack(side=tk.RIGHT, padx=10)
        
        self.start_btn = tk.Button(control_frame, text="启动", command=self.start)
        self.start_btn.pack(pady=5)
        
        self.stop_btn = tk.Button(control_frame, text="停止", command=self.stop)
        self.stop_btn.pack(pady=5)
        
        self.reset_btn = tk.Button(control_frame, text="重置", command=self.reset)
        self.reset_btn.pack(pady=5)

        self.running = False

    def start(self):
        if not self.running:
            self.running = True
            # 随机选择一个物体并设置运动方向
            obj = random.choice(self.objects)
            obj.set_moving(True, angle=random.choice([30, 46, 155, 239]))  # 明确初始方向
            self.moving_objects.append(obj)
            self.update()

    def stop(self):
        self.running = False
        for obj in self.moving_objects:
            obj.set_moving(False)
        self.moving_objects.clear()

    def reset(self):
        self.stop()
        for obj in self.objects:
            obj.x = random.randint(0, WIDTH-1)
            obj.y = random.randint(0, HEIGHT-1)
            obj.color = "black"
            obj.moving = False
            obj.canvas.coords(obj.id, obj.x, obj.y, obj.x+1, obj.y+1)
            obj.canvas.itemconfig(obj.id, fill="black", outline="black")

    def update(self):
        if self.running:
            # 更新所有运动物体位置
            for obj in self.moving_objects:
                obj.move()

            # 感染逻辑：运动物体周围2像素内的静止物体开始运动
            new_movers = []
            for obj in self.objects:
                if not obj.moving:
                    for mover in self.moving_objects:
                        dx = abs(obj.x - mover.x)
                        dy = abs(obj.y - mover.y)
                        if dx <= MIN_DISTANCE and dy <= MIN_DISTANCE:
                            # 设置相同方向并保持距离
                            obj.set_moving(True, angle=mover.angle)
                            new_movers.append(obj)
                            break
            self.moving_objects.extend(new_movers)

            # 维持2像素距离
            for i in range(len(self.moving_objects)):
                for j in range(i+1, len(self.moving_objects)):
                    obj1 = self.moving_objects[i]
                    obj2 = self.moving_objects[j]
                    distance = math.hypot(obj1.x - obj2.x, obj1.y - obj2.y)
                    if distance < MIN_DISTANCE:
                        # 调整obj2的位置
                        angle = math.atan2(obj2.y - obj1.y, obj2.x - obj1.x)
                        obj2.x = int(obj1.x + (MIN_DISTANCE+1) * math.cos(angle))
                        obj2.y = int(obj1.y + (MIN_DISTANCE+1) * math.sin(angle))
                        obj2.canvas.coords(obj2.id, obj2.x, obj2.y, obj2.x+1, obj2.y+1)

            self.root.after(50, self.update)

if __name__ == "__main__":
    root = tk.Tk()
    root.title("群体行为演化模拟")
    sim = EvolutionSimulation(root)
    root.mainloop()